from __future__ import print_function
import tensorflow
import numpy as np
from keras import losses
from pandas import DataFrame
from utils import getData, getLSTM

dataPath = '../data/merge_labeled_data.csv'
saveAdvPath = './data/merge_adv_data.csv'
modelPath = './model/LSTM.hdf5'


# FGSM对抗攻击,产生扰动
def create_adversarial_pattern(input_feature, input_label):
    # 获得网络模型
    model = getLSTM()

    model.load_weights(modelPath)
    model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])

    with tensorflow.GradientTape() as tape:
        input_feature = tensorflow.convert_to_tensor(input_feature, tensorflow.float32, name='input_feature')
        tape.watch(input_feature)
        prediction = model(input_feature)
        # print(prediction)
        loss = losses.binary_crossentropy(input_label, prediction)
        # print(prediction, input_label, loss)

    # 计算导数
    gradient = tape.gradient(loss, input_feature)
    # 扰动值
    signed_grad = (tensorflow.sign(gradient)).numpy().reshape(input_feature.shape[0], input_feature.shape[2])

    return signed_grad


# 生成对抗样本
def FGSM():
    feature, label = getData(dataPath)

    # 扰动值
    eps = 0.1
    # 生成对抗攻击
    adv_pattern = create_adversarial_pattern(feature, label)
    attack = feature.reshape(feature.shape[0], feature.shape[2]) + adv_pattern * eps
    attack = np.clip(attack, .0, 1.)

    adv_attack = np.array(attack).reshape((len(feature), feature.shape[2]))
    print(adv_attack.shape)

    # 融合特征和标签
    adv_attack = np.concatenate((label, adv_attack), axis=1)
    # 保存对抗样本
    data = DataFrame(adv_attack)

    header = ['label', 'diff_ip_num', 'port_entropy', 'all_pkt_num', 'ack_num', 'ack_rate', 'syn_num', 'syn_rate',
              'psh_num', 'psh_rate', 'all_pkt_bytes', 'avg_pkt_bytes', 'one_quarter', 'two_quarter', 'three_quarter']

    data.to_csv(saveAdvPath, header=header, index=False)
    print('成功')


if __name__ == '__main__':
    FGSM()
